import torch
import torch.nn as nn

# ResNet18的基础块：两个卷积层，一个残差连接
class Resnet18BasicBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride1=1, stride2=1, downsample=False):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride1, padding=1, bias=False)
        self.batchnorm1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride2, padding=1, bias=False)
        self.batchnorm2 = nn.BatchNorm2d(out_channels)
        if downsample:
            self.downsample = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride1, bias=False),
                nn.BatchNorm2d(out_channels)
            )
        else:
            self.downsample = None

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.batchnorm1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.batchnorm2(out)
        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

# ResNet18的一个层：两个基础块
class Resnet18Layer(nn.Module):
    def __init__(self,in_channels, out_channels,stride = 1, downsample=False):
        super().__init__()
        self.layer = nn.Sequential(
            Resnet18BasicBlock(in_channels, out_channels,stride1 = stride, downsample=downsample),
            Resnet18BasicBlock(out_channels, out_channels)
        )

    def forward(self, x):
        return self.layer(x)

class ResNet18(nn.Module):
    def __init__(self, initweight = False):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.batchnorm1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        self.layer1 = Resnet18Layer(64, 64)
        self.layer2 = Resnet18Layer(64, 128, stride=2, downsample=True)

        self.layer3 = Resnet18Layer(128, 256, stride=2, downsample=True)
        self.layer4 = Resnet18Layer(256, 512, stride=2, downsample=True)


    def forward(self, x):
        '''
        输出为一个字典, key为X0, X1, X2, X3_sm, X4_sm, X3_im, X4_im
        X0 - X4:分别是五个层的输出
        X3/4_sm和X3/4_im是两个分支的输出
        '''
        x = self.conv1(x)
        x = self.batchnorm1(x)
        x = self.relu(x)
        x0 = self.maxpool(x)

        x1 = self.layer1(x0)
        x2 = self.layer2(x1)
        

        x3 = self.layer3(x2)
        x4 = self.layer4(x3)

        return {"X0": x0, "X1": x1, "X2": x2, "X3": x3, "X4": x4}

if __name__ == "__main__":
    model = ResNet18(True)
    input = torch.randn(32, 3, 156, 156)
    output = model(input)
    for key in output.keys():
        print(key, output[key].shape)